import os
import cv2
import numpy as np
from flask import Flask, request, jsonify, send_from_directory
from werkzeug.utils import secure_filename
import math
from img_processing import ImgProcessing

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads/'
app.config['PROCESSED_FOLDER'] = 'static/processed/'
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'webp'}

# 创建必要的文件夹
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['PROCESSED_FOLDER'], exist_ok=True)


def allowed_file(filename):
    """检查文件是否为允许的类型"""
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


def preprocess_image(image):
    """图像预处理：灰度化、阈值分割、降噪"""
    # 转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 高斯模糊降噪
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # 阈值分割
    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 形态学操作去除噪声
    kernel = np.ones((3, 3), np.uint8)
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)

    return opening, gray


def correct_perspective(image, preprocessed):
    """校正透视变形和倾斜"""
    # 寻找轮廓
    contours, _ = cv2.findContours(preprocessed.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if not contours:
        return image  # 如果没有找到轮廓，返回原图

    # 找到最大的轮廓（假设是零件本身）
    largest_contour = max(contours, key=cv2.contourArea)

    # 近似多边形
    perimeter = cv2.arcLength(largest_contour, True)
    approx = cv2.approxPolyDP(largest_contour, 0.02 * perimeter, True)

    # 如果找到四边形，进行透视变换
    if len(approx) == 4:
        # 排序顶点
        approx = approx.reshape(4, 2)
        rect = np.zeros((4, 2), dtype="float32")

        # 按总和排序（左上角和右下角）
        s = approx.sum(axis=1)
        rect[0] = approx[np.argmin(s)]
        rect[2] = approx[np.argmax(s)]

        # 按差值排序（右上角和左下角）
        diff = np.diff(approx, axis=1)
        rect[1] = approx[np.argmin(diff)]
        rect[3] = approx[np.argmax(diff)]

        # 计算目标宽度和高度
        (tl, tr, br, bl) = rect
        widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
        widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
        maxWidth = max(int(widthA), int(widthB))

        heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
        heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
        maxHeight = max(int(heightA), int(heightB))

        # 目标点
        dst = np.array([
            [0, 0],
            [maxWidth - 1, 0],
            [maxWidth - 1, maxHeight - 1],
            [0, maxHeight - 1]], dtype="float32")

        # 计算透视变换矩阵并应用
        M = cv2.getPerspectiveTransform(rect, dst)
        warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))

        return warped
    else:
        # 如果不是四边形，尝试旋转校正（基于最小外接矩形）
        rect = cv2.minAreaRect(largest_contour)
        angle = rect[-1]

        # 调整角度
        if angle < -45:
            angle += 90

        # 获取旋转矩阵
        (h, w) = image.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)

        # 执行旋转
        rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)

        return rotated


def detect_holes(image):
    """检测圆形孔洞并返回中心坐标"""
    # 转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 高斯模糊
    gray_blurred = cv2.GaussianBlur(gray, (9, 9), 2)

    # 霍夫圆检测
    circles = cv2.HoughCircles(
        gray_blurred,  # 输入图像（灰度图，建议先模糊处理）
        method=cv2.HOUGH_GRADIENT,  # 检测方法
        dp=1,  # 累加器分辨率与图像分辨率的比值
        minDist=100,  # 检测到的圆之间的最小距离
        param1=200,  # 边缘检测的高阈值（Canny边缘检测）
        param2=60,  # 累加器阈值（决定检测到的圆的数量和质量）
        minRadius=10,  # 圆的最小半径
        maxRadius=90  # 圆的最大半径
    )

    # 处理检测结果
    hole_centers = []
    if circles is not None:
        # 转换为整数坐标
        circles = np.uint16(np.around(circles))
        for circle in circles[0, :]:
            center = (circle[0], circle[1])
            hole_centers.append(center)

            # 在图像上标记圆心
            cv2.circle(image, center, 5, (0, 255, 0), -1)
            # 标记圆
            cv2.circle(image, center, circle[2], (0, 0, 255), 2)

    return image, hole_centers


# def is_square(points, tolerance=10):
#     """检查四个点是否构成正方形（允许一定误差）"""
#     if len(points) != 4:
#         return False, "点的数量不是4个"
#
#     # 计算所有点之间的距离
#     distances = []
#     for i in range(4):
#         for j in range(i + 1, 4):
#             dist = math.hypot(points[i][0] - points[j][0], points[i][1] - points[j][1])
#             distances.append(dist)
#
#     # 排序距离
#     distances.sort()
#
#     # 正方形应该有4条边和2条对角线
#     # 前4个应该是边长，后2个应该是对角线
#     side_distances = distances[:4]
#     diagonal_distances = distances[4:]
#
#     # 检查四条边是否大致相等
#     avg_side = sum(side_distances) / 4
#     # for d in side_distances:
#     #     if abs(d - avg_side) > tolerance:
#     #         return False, f"边的长度差异超过阈值，最大差异: {max(abs(d - avg_side) for d in side_distances):.2f}px"
#
#     # 检查两条对角线是否大致相等
#     if len(diagonal_distances) != 2:
#         return False, "对角线数量不正确"
#
#     if abs(diagonal_distances[0] - diagonal_distances[1]) > tolerance:
#         return False, f"对角线长度差异超过阈值: {abs(diagonal_distances[0] - diagonal_distances[1]):.2f}px"
#
#     # 检查对角线是否约为边长的√2倍
#     expected_diagonal = avg_side * math.sqrt(2)
#     avg_diagonal = sum(diagonal_distances) / 2
#     if abs(avg_diagonal - expected_diagonal) > tolerance * math.sqrt(2):
#         return False, f"对角线与边长比例不符合正方形特性，差异: {abs(avg_diagonal - expected_diagonal):.2f}px"
#
#     return True, "符合正方形特征"

def is_square(points, tolerance=100):
    """检查四个点是否构成长方形（允许一定误差）"""
    if len(points) != 4:
        return False, "点的数量不是4个"

    # 计算所有点之间的距离
    points_np = np.array(points, dtype=np.float64)
    distances = []
    point_pairs = []  # 保存距离对应的点对，用于后续分析
    for i in range(4):
        for j in range(i + 1, 4):
            p1 = points_np[i]
            p2 = points_np[j]


            dist = np.linalg.norm(p1 - p2)
            distances.append(dist)
            point_pairs.append((i, j, dist))

    # 排序距离和对应的点对
    sorted_pairs = sorted(point_pairs, key=lambda x: x[2])
    distances.sort()

    # 长方形有4条边（2对相等的边）和2条对角线
    # 前4个为边，后2个为对角线
    sides = distances[:4]
    diagonals = distances[4:]

    # 检查边是否形成两对相等的边（长和宽）
    # 允许一定误差范围内的相等
    def are_equal(a, b, tol):
        return abs(a - b) <= tol

    # 检查边是否能组成两组对边
    if not (are_equal(sides[0], sides[1], tolerance) and are_equal(sides[2], sides[3], tolerance)) and \
            not (are_equal(sides[0], sides[2], tolerance) and are_equal(sides[1], sides[3], tolerance)) and \
            not (are_equal(sides[0], sides[3], tolerance) and are_equal(sides[1], sides[2], tolerance)):
        return False, f"边不能形成两组相等的对边，边长为: {[round(d, 2) for d in sides]}"

    # 检查两条对角线是否相等
    if len(diagonals) != 2:
        return False, "对角线数量不正确"

    if not are_equal(diagonals[0], diagonals[1], tolerance):
        return False, f"对角线长度不相等，差异: {abs(diagonals[0] - diagonals[1]):.2f}px"

    # 验证四个角是否为直角（通过向量点积判断垂直）
    # 先找到四个点的正确顺序（按顺时针或逆时针排列）
    # 计算各点的中心
    center = (sum(p[0] for p in points) / 4, sum(p[1] for p in points) / 4)

    # 计算各点相对于中心的角度，用于排序
    def get_angle(point):
        return math.atan2(point[1] - center[1], point[0] - center[0])

    # 按角度排序点，得到顺时针/逆时针顺序
    sorted_points = sorted(points, key=get_angle)

    # 计算相邻边的向量并检查是否垂直（点积接近0）
    for i in range(4):
        p1 = sorted_points[i]
        p2 = sorted_points[(i + 1) % 4]
        p3 = sorted_points[(i + 2) % 4]

        # 向量1: p2 - p1
        v1 = (p2[0] - p1[0], p2[1] - p1[1])
        # 向量2: p3 - p2
        v2 = (p3[0] - p2[0], p3[1] - p2[1])

        # 点积
        dot_product = v1[0] * v2[0] + v1[1] * v2[1]

        # 向量长度乘积
        len_v1 = math.hypot(v1[0], v1[1])
        len_v2 = math.hypot(v2[0], v2[1])

        if len_v1 == 0 or len_v2 == 0:
            return False, "存在重合的点"

        # 计算余弦值（接近0表示垂直）
        cos_theta = abs(dot_product) / (len_v1 * len_v2)

        # 允许一定误差（余弦值小于0.1视为接近直角）
        if cos_theta > 0.1:
            return False, f"角不是直角，余弦值: {cos_theta:.2f}"

    return True, "符合长方形特征"

@app.route('/upload', methods=['POST'])
def upload_file():
    """处理文件上传并进行检测"""
    # 检查是否有文件上传
    if 'file' not in request.files:
        return jsonify({'error': '没有文件上传'}), 400

    file = request.files['file']

    # 检查文件名是否为空
    if file.filename == '':
        return jsonify({'error': '未选择文件'}), 400

    # 检查文件类型
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        upload_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(upload_path)

        # 读取图像
        image = cv2.imread(upload_path)
        if image is None:
            return jsonify({'error': '无法读取图像文件'}), 400

        # 预处理图像
        # preprocessed, gray = preprocess_image(image)
        #
        # # 校正透视和倾斜
        # corrected = correct_perspective(image.copy(), preprocessed)
        #
        # # 检测孔洞
        # processed_image, hole_centers = detect_holes(corrected)
        #
        # # 保存处理后的图像
        processed_filename = f"processed_{filename}"
        processed_path = os.path.join(app.config['PROCESSED_FOLDER'], processed_filename)
        # cv2.imwrite(processed_path, processed_image)
        #
        imgpro = ImgProcessing(image)
        # dis = imgpro.distance()
        img_end = imgpro.img_show(processed_path)

        # # 检查是否为正方形
        is_square_result, message = is_square(imgpro.img_circle) if len(imgpro.img_circle) == 4 else (False,
                                                                                            "检测到的孔数量不是4个")



        # 准备返回结果
        result = {
            'original_filename': filename,
            'processed_filename': processed_filename,
            'hole_count': len(imgpro.img_circle),
            'hole_centers': [{'x': int(x), 'y': int(y)} for (x, y) in imgpro.img_circle],
            'is_square': is_square_result,
            'message': message
        }

        return jsonify(result)

    return jsonify({'error': '不支持的文件类型'}), 400


@app.route('/uploads/<filename>')
def uploaded_file(filename):
    """提供上传的原始文件"""
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)


@app.route('/processed/<filename>')
def processed_file(filename):
    """提供处理后的文件"""
    return send_from_directory(app.config['PROCESSED_FOLDER'], filename)


@app.route('/')
def index():
    """返回前端页面"""
    return send_from_directory('.', 'index.html')


if __name__ == '__main__':
    app.run(debug=True)
